import numpy as np
import matplotlib.pyplot as plt


def SaltAndPepperNoise(img, percentage=0.05):
    """
    为图像添加椒盐噪声，由定义可知，椒盐噪声是随机分布的黑、白像素点
    :param img: 数据类型numpy ndarray，范围0-255 uint8，图像形状H*W*3，顺序RGB，或H*W灰度图
    :param percentage: 添加噪点的百分比
    :return:
    """
    img_dst = img.copy()  # 保持原图像不变
    num_noise_point = int(percentage * img.shape[0] * img.shape[1])  # 计算噪点数目
    for i in range(num_noise_point):
        x = np.random.randint(0, img.shape[0] - 1)
        y = np.random.randint(0, img.shape[1] - 1)
        p_black = np.random.randint(0, 1)  # 随机生成白点或黑点
        # 对三通道图和单通道灰度图分别处理
        if len(img) == 3:
            img_dst[x, y, :] = 0 if p_black == 0 else 255
        else:
            img_dst[x, y] = 0 if p_black == 0 else 255
    return img_dst


def GaussianNoise(img, scale=5):
    """
    为图像添加零均值高斯噪声
    :param img: 数据类型numpy ndarray，范围0-255 uint8，图像形状H*W*3，顺序RGB，或H*W灰度图
    :param scale: 噪点分布的方差
    :return:
    """
    if len(img.shape) == 2:
        h, w = img.shape
    else:
        h, w, _ = img.shape
    img_dst = img.copy()
    if len(img.shape) == 3:
        for x in range(h):
            for y in range(w):
                s = np.random.normal(loc=0, scale=scale, size=3)
                img_dst[x, y, 0] = img[x, y, 0] + s[0]
                img_dst[x, y, 1] = img[x, y, 1] + s[1]
                img_dst[x, y, 2] = img[x, y, 2] + s[2]
    else:
        for x in range(h):
            for y in range(w):
                s = np.random.normal(loc=0, scale=scale, size=1)
                img_dst[x, y] = img[x, y] + s
    np.clip(img_dst, 0, 255)
    return img_dst


def median_filter(img, kernal_size=3):
    """
    中值滤波器
    :param img: 数据类型numpy ndarray，范围0-255 uint8，图像形状H*W*3，顺序RGB，或H*W灰度图
    :param kernal_size: 滤波器核大小
    :return:
    """
    pad = kernal_size // 2
    img_dst = img.copy()
    if len(img.shape) == 2:
        h, w = img.shape
        # 边框补零
        img_pad = np.pad(img, ((pad, pad), (pad, pad)))
        # 中值滤波
        for y in range(0, h):
            for x in range(0, w):
                img_dst[y, x] = np.median(img_pad[y:y + kernal_size, x:x + kernal_size])
    else:
        h, w, c = img.shape
        # 边框补零
        img_pad = np.pad(img, ((pad, pad), (pad, pad), (0, 0)))
        # 中值滤波
        for y in range(0, h):
            for x in range(0, w):
                for ch in range(c):
                    img_dst[y, x, ch] = np.median(img_pad[y:y + kernal_size, x:x + kernal_size, ch])
    np.clip(img_dst, 0, 255)
    img_dst = img_dst.astype(np.uint8)
    # plt.figure()
    # plt.imshow(img_dst)
    # plt.show()
    return img_dst


def average_filter(img, kernal_size=3):
    """
    均值滤波器
    :param img: 数据类型numpy ndarray，范围0-255 uint8，图像形状H*W*3，顺序RGB，或H*W灰度图
    :param kernal_size: 滤波器核大小
    :return:
    """
    pad = kernal_size // 2
    img_dst = img.copy()
    if len(img.shape) == 2:
        h, w = img.shape
        # 边框补零
        img_pad = np.pad(img, ((pad, pad), (pad, pad)))
        # 中值滤波
        for y in range(0, h):
            for x in range(0, w):
                img_dst[y, x] = np.mean(img_pad[y:y + kernal_size, x:x + kernal_size])
    else:
        h, w, c = img.shape
        # 边框补零
        img_pad = np.pad(img, ((pad, pad), (pad, pad), (0, 0)))
        # 中值滤波
        for y in range(0, h):
            for x in range(0, w):
                for ch in range(c):
                    img_dst[y, x, ch] = np.mean(img_pad[y:y + kernal_size, x:x + kernal_size, ch])
    np.clip(img_dst, 0, 255)
    img_dst = img_dst.astype(np.uint8)
    # plt.figure()
    # plt.imshow(img_dst)
    # plt.show()
    return img_dst


if __name__ == '__main__':
    image0 = plt.imread("1.jpg")
    # print(image0.dtype)
    # print(np.max(image0))
    plt.subplot(3, 3, 1)
    plt.title("original image")
    plt.imshow(image0)

    SAP_image = SaltAndPepperNoise(image0, 0.05)
    plt.subplot(3, 3, 2)
    plt.title("SAP_image")
    plt.imshow(SAP_image)

    G_image = GaussianNoise(image0)
    plt.subplot(3, 3, 3)
    plt.title("G_image")
    plt.imshow(G_image)

    m_img1 = median_filter(SAP_image)
    plt.subplot(3, 3, 5)
    plt.title("median_SP")
    plt.imshow(m_img1)

    a_img1 = average_filter(SAP_image)
    plt.subplot(3, 3, 6)
    plt.title("average_SP")
    plt.imshow(a_img1)

    m_img2 = median_filter(G_image)
    plt.subplot(3, 3, 8)
    plt.title("median_G")
    plt.imshow(m_img2)

    a_img2 = average_filter(G_image)
    plt.subplot(3, 3, 9)
    plt.title("average_G")
    plt.imshow(a_img2)

    plt.show()
